#!/usr/bin/env python
# -*- encoding: UTF-8 -*-
from __future__ import print_function

import numpy as np
import cv2
import copy

# 机器人工具箱
from robot import *
# Nao 的参数
from nao_defines import *

def RotXYZMatrix(xAngle,yAngle,zAngle):
    Rx = rotx(xAngle)
    Ry = roty(yAngle)
    Rz = rotz(zAngle)
    R = Rx*Ry*Rz
    return r2t(R)

def RotZYXMatrix(zAngle,yAngle,xAngle):
    Rx = rotx(xAngle)
    Ry = roty(yAngle)
    Rz = rotz(zAngle)
    R = Rz*Ry*Rx
    return r2t(R)

def DH(a, alpha, d, theta):
    Taf = [[cos(theta),            -sin(theta),            0,              a],
    [sin(theta)*cos(alpha), cos(theta)*cos(alpha),  -sin(alpha),    -sin(alpha)*d],
    [sin(theta)*sin(alpha), cos(theta)*sin(alpha),  cos(alpha),     cos(alpha)*d],
    [0,                     0,                      0,              1]]
    return mat(Taf)


class LeftHand(object):
    def __init__(self):
        self.Base = eye(4)
        self.Base[1,3] = ShoulderOffsetY
        self.Base[2,3] = ShoulderOffsetZ
        self.T1 = lambda theta: DH(0,-pi/2,0,theta)        # LShoulderPitch
        self.T2 = lambda theta: DH(0,pi/2,0,theta+pi/2)    # LShoulderRoll
        self.T3 = lambda theta: DH(ElbowOffsetY,pi/2,UpperArmLength,theta) # LElbowYaw
        self.T4 = lambda theta: DH(0,-pi/2,0,theta)        # LElbowRoll
        self.T5 = lambda theta: DH(0,pi/2,0,theta)         # LWristYaw
        self.Tool = eye(4)
        self.Tool[0,3] = LowerArmLength + HandOffsetX
        self.Tool[2,3] = -HandOffsetZ
        self.Tool = RotXYZMatrix(-pi/2,0,-pi/2)*self.Tool

class RightHand(object):
    def __init__(self):
        self.Base = eye(4)
        self.Base[1,3] = -ShoulderOffsetY # 与LArm不同
        self.Base[2,3] = ShoulderOffsetZ
        self.T1 = lambda theta: DH(0,-pi/2,0,theta)       # RShoulderPitch
        self.T2 = lambda theta: DH(0,pi/2,0,theta+pi/2)   # RShoulderRoll
        self.T3 = lambda theta: DH(-ElbowOffsetY,pi/2,UpperArmLength,theta) # RElbowYaw 与 LArm 不同
        self.T4 = lambda theta: DH(0,-pi/2,0,theta)       # RElbowRoll
        self.T5 = lambda theta: DH(0,pi/2,0,theta)        # RWristYaw
        self.Tool = eye(4)
        self.Tool[0,3] = LowerArmLength + HandOffsetX
        self.Tool[2,3] = -HandOffsetZ
        self.Tool = RotXYZMatrix(-pi/2,0,-pi/2)*self.Tool

class LeftLeg(object):
    def __init__(self):
        self.Base = eye(4)
        self.Base[1,3] = HipOffsetY
        self.Base[2,3] = -HipOffsetZ
        self.T1 = lambda theta: DH(0,-3*pi/4,0,theta-pi/2)  # LHipYawPitch
        self.T2 = lambda theta: DH(0,-pi/2,0,theta+pi/4)    # LHipRoll
        self.T3 = lambda theta: DH(0,pi/2,0,theta)          # LHipPitch
        self.T4 = lambda theta: DH(-ThighLength,0,0,theta)  # LKneePitch
        self.T5 = lambda theta: DH(-TibiaLength,0,0,theta)  # LAnklePitch
        self.T6 = lambda theta: DH(0,-pi/2,0,theta)         # LAnkleRoll
        self.Tool = eye(4)
        self.Tool[2,3] = -FootHeight
        self.Tool = RotZYXMatrix(pi,-pi/2,0) * self.Tool

class RightLeg(object):
    def __init__(self):
        self.Base = eye(4)
        self.Base[1,3] = -HipOffsetY # 与 LeftLeg 不同
        self.Base[0,3] = -HipOffsetZ
        self.T1 = lambda theta: DH(0,-pi/4,0,theta-pi/2)   # RHipYawPitch 与 LeftLeg 不同
        self.T2 = lambda theta: DH(0,-pi/2,0,theta-pi/4)   # RHipRoll 与 LeftLeg 不同
        self.T3 = lambda theta: DH(0,pi/2,0,theta)         # RHipPitch
        self.T4 = lambda theta: DH(-ThighLength,0,0,theta) # RKneePitch
        self.T5 = lambda theta: DH(-TibiaLength,0,0,theta) # RAnklePitch
        self.T6 = lambda theta: DH(0,-pi/2,0,theta)        # RAnkleRoll
        self.Tool = eye(4)
        self.Tool[2,3] = -FootHeight
        self.Tool = RotZYXMatrix(pi,-pi/2,0) * self.Tool

class Head(object):
    def __init__(self):
        self.Base = eye(4)
        self.Base[2,3] = NeckOffsetZ
        self.T1 = lambda theta: DH(0,0,0,theta)            # HeadYaw
        self.T2 = lambda theta: DH(0,-pi/2,0,theta-pi/2)   # HeadPitch
        ## ???
        self.ToolDownCamera = eye(4)
        self.ToolDownCamera[0,3] = CameraBotomX
        self.ToolDownCamera[2,3] = CameraBotomZ
        self.ToolDownCamera = RotXYZMatrix(pi/2,pi/2,0) *  self.ToolDownCamera

# 单例修饰器
def singleton(cls, *args, **kw):
    instances = {}
    def _singleton():
        if cls not in instances:
            instances[cls] = cls(*args, **kw)
        return instances[cls]
    return _singleton

@singleton
class Kinematics(object):
    def __init__(self):
        self.Head = Head()
        self.LeftHand = LeftHand()
        self.RightHand = RightHand()
        self.LeftLeg = LeftLeg()
        self.RightLeg = RightLeg()

class Nao(Kinematics):
    def __init__(self):
        super(Nao, self).__init__(self)
        pass

##### 正运动学

def tf2pos(tf):
    wz,wy,wx = array(tr2rpy(tf)).flatten()
    pos = zeros((6,1))
    pos[0:3,0] = tf[0:3,3].transpose()
    pos[3:6,0] = [wx,wy,wz]
    return pos

def pos2tf(pos):
    x,y,z,wx,wy,wz = array(pos).flatten()
    tf = eye(4)
    tf[0:3,3] = [x,y,z]
    tf[:3,:3] = rpy2tr(wz,wy,wx)[:3,:3]
    return tf

def fLeftHand(theta1,theta2,theta3,theta4,theta5,theta6=0):
    """
    NAO H25 (v 4.0, v5.0) Left Hand forward kinematics
    :param theta1:
    :param theta2:
    :param theta3:
    :param theta4:
    :param theta5:
    :return:
    """
    kin = Kinematics()
    Base = kin.LeftHand.Base
    T1 = kin.LeftHand.T1(theta1)
    T2 = kin.LeftHand.T2(theta2)
    T3 = kin.LeftHand.T3(theta3)
    T4 = kin.LeftHand.T4(theta4)
    T5 = kin.LeftHand.T5(theta5)
    Tool = kin.LeftHand.Tool
    Tend = Base*T1*T2*T3*T4*T5*Tool
    pos = tf2pos(Tend)
    return Tend, pos

def fRightHand(theta1,theta2,theta3,theta4,theta5,theta6=0):
    """
    NAO H25 (v 4.0, v5.0) Right Hand forward kinematics
    :param theta1:
    :param theta2:
    :param theta3:
    :param theta4:
    :param theta5:
    :return:
    """
    kin = Kinematics()
    Base = kin.RightHand.Base
    T1 = kin.RightHand.T1(theta1)
    T2 = kin.RightHand.T2(theta2)
    T3 = kin.RightHand.T3(theta3)
    T4 = kin.RightHand.T4(theta4)
    T5 = kin.RightHand.T5(theta5)
    Tool = kin.RightHand.Tool
    Tend = Base*T1*T2*T3*T4*T5*Tool
    pos = tf2pos(Tend)
    return Tend, pos

def fLeftLeg(theta1,theta2,theta3,theta4,theta5,theta6):
    """
    NAO H25 (v 4.0, v5.0) Left Leg forward kinematics
    :param theta1:
    :param theta2:
    :param theta3:
    :param theta4:
    :param theta5:
    :return:
    """
    kin = Kinematics()
    Base = kin.LeftLeg.Base
    T1 = kin.LeftLeg.T1(theta1)
    T2 = kin.LeftLeg.T2(theta2)
    T3 = kin.LeftLeg.T3(theta3)
    T4 = kin.LeftLeg.T4(theta4)
    T5 = kin.LeftLeg.T5(theta5)
    T6 = kin.LeftLeg.T6(theta6)
    Tool = kin.LeftLeg.Tool
    Tend = Base*T1*T2*T3*T4*T5*T6*Tool
    pos = tf2pos(Tend)
    return Tend, pos

def fRightLeg(theta1,theta2,theta3,theta4,theta5,theta6):
    """
    NAO H25 (v 4.0, v5.0) Right Leg forward kinematics
    :param theta1:
    :param theta2:
    :param theta3:
    :param theta4:
    :param theta5:
    :return:
    """
    kin = Kinematics()
    Base = kin.RightLeg.Base
    T1 = kin.RightLeg.T1(theta1)
    T2 = kin.RightLeg.T2(theta2)
    T3 = kin.RightLeg.T3(theta3)
    T4 = kin.RightLeg.T4(theta4)
    T5 = kin.RightLeg.T5(theta5)
    T6 = kin.RightLeg.T6(theta6)
    Tool = kin.RightLeg.Tool
    Tend = Base*T1*T2*T3*T4*T5*T6*Tool
    pos = tf2pos(Tend)
    return Tend, pos

def fDownCamera(theta1,theta2):
    """
    NAO H25 (v 4.0, v5.0) Down Camera forward kinematics
    :param theta1:
    :param theta2:
    :param theta3:
    :param theta4:
    :param theta5:
    :return:
    """
    kin = Kinematics()
    Base = kin.Head.Base
    T1 = kin.Head.T1(theta1)
    T2 = kin.Head.T2(theta2)
    Tool = kin.Head.DownCamera
    Tend = Base*T1*T2*Tool
    pos = tf2pos(Tend)
    return Tend, pos

def fUpCamera(theta1,theta2):
    """
    NAO H25 (v 4.0, v5.0) Up Camera forward kinematics
    :param theta1:
    :param theta2:
    :param theta3:
    :param theta4:
    :param theta5:
    :return:
    """
    Tend, pos = fDownCamera(theta1,theta2) # 未完成
    return Tend, pos

##### 逆运动学
from numpy.linalg import *

def iLeftHand(Tend,theta1=None):
    kin = Kinematics()
    Base = kin.LeftHand.Base
    Tool = kin.LeftHand.Tool
    T15 = inv(Tool*inv(Tend)*Base)
    theta1 = atan(-T15[2,3]/T15[0,3]) if theta1 is None else theta1
    T25 = inv(kin.LeftHand.T1(theta1))*T15
    l1 = UpperArmLength
    l2 = ElbowOffsetY
    r14 = T25[0,3]
    r34 = T25[2,3]
    theta2 = atan2(l1*r34-l2*r14,l1*r14+l2*r34)
    T35 = inv(kin.LeftHand.T2(theta2))*T25
    theta3 = atan2(-T35[2,2],-T35[0,2])
    T45 =inv(kin.LeftHand.T3(theta3))*T35
    theta4 = atan2(T45[0,2],T45[2,2])
    theta5 = atan2(T45[1,0],T45[1,1])
    return theta1,theta2,theta3,theta4,theta5,0

def iRightHand(Tend,theta1=None):
    kin = Kinematics()
    Base = kin.RightHand.Base
    Tool = kin.RightHand.Tool
    T15 = inv(Tool*inv(Tend)*Base)
    theta1 = atan(-T15[2,3]/T15[0,3]) if theta1 is None else theta1
    T25 = inv(kin.RightHand.T1(theta1))*T15
    l1 = UpperArmLength
    l2 = ElbowOffsetY
    r14 = T25[0,3]
    r34 = T25[2,3]
    theta2 = atan2(l1*r34+l2*r14,l1*r14-l2*r34) # 与 Left 不同
    T35 = inv(kin.RightHand.T2(theta2))*T25
    theta3 = atan2(T35[2,2],T35[0,2]) # 与 Left 不同
    T45 = inv(kin.RightHand.T3(theta3))*T35
    theta4 = atan2(T45[0,2],T45[2,2])
    theta5 = atan2(T45[1,0],T45[1,1])
    return theta1,theta2,theta3,theta4,theta5,0


##### 绕定轴转动

def get_p_w_theta(Tf):
    Tf = np.array(Tf)
    p = Tf[:3,3]
    w = cv2.Rodrigues(Tf[:3,:3])[0][:,0] #
    th = np.linalg.norm(w)
    w = w/th
    return p,w,th

def getTf(p,w,theta):
    Tf = np.eye(4)
    Tf[:3,3] = p
    Tf[:3,:3] = cv2.Rodrigues(w*theta)[0]
    return Tf


from scipy.optimize import minimize

l1 = UpperArmLength
l2 = ElbowOffsetY
l3 = ShoulderOffsetY
l4 = ShoulderOffsetZ
l5 = HandOffsetX
l6 = HandOffsetZ
l7 = LowerArmLength

def get_theta1_from_theta(theta, w, p):
    px,py,pz = p
    rx,ry,rz = w
    th = theta
    A1 = -l6*rx*rz + l5*rx**2 + l7*rx**2 - l5 - l7
    A2 = l6*ry
    A3 = l6*rx*rz + px - (l5+l7)*rx**2
    a = cos(th)*A1 + sin(th)*A2 + A3
    B1 = -l6*rz**2 + l6 + l5*rz*rz + l7*rx*rz
    B2 = l5*ry + l7*ry
    B3 = -l4 + l6 + pz - rx*rz*(l5+l7)
    b = B1*cos(th) + B2*sin(th) + B3
    theta1 = atan(-b/a)
    return theta1

from scipy.optimize import minimize
def my_ik(fkfunc, ikfunc, Tend):
    """
    绕w轴转动, 并沿着该转动轴上下移动 dp, 移动量不超过 20 mm
    :param fkfunc: 正运动学函数, fLeftHand 或 fRightHand
    :param ikfunc: 逆运动学函数, iLeftHand 或 iRightHand
    :param Tend: 目标末端位姿
    :return: 误差最小的解
    """
    p0,w,th0 = get_p_w_theta(Tend)
    get_Tend = lambda theta,dp: getTf(p0+dp*w,w,theta)
    def func(x):
        theta, dp = x
        newTend = get_Tend(theta,dp)
        try:
            p1, w1, th1 = get_p_w_theta(fkfunc(*ikfunc(newTend))[0])
            return np.linalg.norm(p0+dp*w-p1)
        except:
            return 1000.0

    x0 = [th0, 0]
    bnds = ((0,2*pi),(-20,20))
    res = minimize(func,x0,bounds=bnds,
                   options={
                       'gtol':1e-6,
                       'disp':True
                   })
    new_th,new_dp = res.x
    newTend = get_Tend(new_th,new_dp)
    return ikfunc(newTend)

def my_iLeftHand(Tend):
    return my_ik(fLeftHand,iLeftHand,Tend)

def my_iRightHand(Tend):
    return my_ik(fRightHand,iRightHand,Tend)